import pandas as pd
import numpy as np
import random
import os
from config import dummy_data_dir

# 设置随机种子保证结果可复现
np.random.seed(42)
random.seed(42)

# 1. 准备10个与中国移动有合作的非中国移动企业名称
group_names = [
    "阿里巴巴集团",
    "腾讯控股集团",
    "中国工商银行集团",
    "国家电网公司",
    "华为技术有限公司",
    "中国石油集团",
    "中国建筑集团",
    "顺丰控股集团",
    "新东方教育集团",
    "华润医药集团"
]

# 2. 生成集团用户信息
# 从U00001-U03000中随机选择1000个用户作为企业用户
user_ids = [f"U{str(i).zfill(5)}" for i in range(1, 3001)]
enterprise_users = random.sample(user_ids, 1000)

# 3. 为企业用户分配集团
group_user_mapping = []

# 确定每个集团的用户数量（使用帕累托分布）
group_counts = {}
for i, group in enumerate(group_names):
    # 前20%的企业占据60%的用户（帕累托分布）
    if i < 2:  # 阿里巴巴和腾讯
        num_users = int(len(enterprise_users) * 0.20)
    elif i < 5:  # 工行、国网、华为
        num_users = int(len(enterprise_users) * 0.15)
    else:  # 剩余5个企业
        num_users = int(len(enterprise_users) * 0.15 / 5)
    group_counts[group] = num_users

print(group_counts)
# 调整总用户数
total_allocated = sum([group_counts[group] for group in group_names])
if total_allocated != len(enterprise_users):
    # 调整最后一个企业的用户数
    group_counts[group_names[-1]] += len(enterprise_users) - total_allocated

# 分配用户到集团
start_index = 0
for group in group_names:
    num_users = group_counts[group]
    end_index = start_index + num_users
    
    # 分配用户
    for user in enterprise_users[start_index:end_index]:
        group_user_mapping.append({
            "User ID": user,
            "Group Name": group
            
        })
    
    start_index = end_index

# 4. 创建DataFrame并保存为Excel
df_group = pd.DataFrame(group_user_mapping)
output_file = os.path.join(dummy_data_dir, '集团用户信息表.xlsx')
df_group.to_excel(output_file, index=False)

print(f"集团用户信息生成完成，共{len(df_group)}条记录，已保存至: {output_file}")
print(f"集团数量: {len(group_names)}")
print("\n各集团用户分布:")

for group in group_names:
    count = group_counts[group]
    print(f"{group}: {count}个用户")